import pandas as pd
import os
import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.font_manager as font_manager
from typing import List, Dict, Tuple


class IndustryExpansionRecommender:
    def __init__(self, data_dir: str = "hs300_data"):
        """初始化行业拓展推荐器，加载行业相关度数据并配置可视化"""
        self.data_dir = data_dir
        self.industry_corr = {}  # 行业间相关度映射：{(行业1, 行业2): 相关度分数}
        self.all_industries = []  # 所有有效行业列表
        self.chinese_font = None  # 用于中文显示的字体

        # 初始化流程：字体配置→数据加载
        self.chinese_font = self._get_valid_chinese_font()
        self._load_industry_correlation()
        self._get_all_industries()
        print("✅ 初始化完成！支持行业拓展推荐（相关行业/对冲行业）及可视化")
        print("📌 提示：输入行业时可参考示例（直接复制更准确）：")
        print(f"   示例行业：{self.all_industries}\n")

    # ---------------- 中文字体配置 ----------------
    def _get_valid_chinese_font(self) -> str:
        """检测并返回可用的中文字体（确保图表中文显示）"""
        candidate_fonts = [
            "SimHei", "Microsoft YaHei", "PingFang SC", 
            "Heiti SC", "WenQuanYi Zen Hei"
        ]
        
        # 检测系统自带字体
        available_fonts = set(f.name for f in font_manager.fontManager.ttflist)
        for font in candidate_fonts:
            if font in available_fonts:
                print(f"✅ 检测到可用中文字体：{font}（用于图表渲染）")
                return font
        
        # 加载本地备用字体（若系统无自带字体）
        custom_font_path = os.path.join(self.data_dir, "fonts", "SimHei.ttf")
        if os.path.exists(custom_font_path):
            font_manager.fontManager.addfont(custom_font_path)
            font_manager._rebuild()
            if "SimHei" in [f.name for f in font_manager.fontManager.ttflist]:
                print(f"✅ 已加载本地字体：{custom_font_path}")
                return "SimHei"
        
        # 无可用字体时提示解决方案
        raise RuntimeError(
            "❌ 未找到可用中文字体！请按以下步骤解决：\n"
            "1. 下载 SimHei.ttf 字体文件\n"
            "2. 创建目录：hs300_data/fonts\n"
            "3. 将 SimHei.ttf 放入该目录后重试"
        )

    # ---------------- 数据加载 ----------------
    def _load_industry_correlation(self):
        """加载行业相关度数据（从industry_correlation.csv读取）"""
        corr_file = os.path.join(self.data_dir, "industry_correlation.csv")
        if not os.path.exists(corr_file):
            raise FileNotFoundError(
                f"❌ 缺少行业相关度数据文件：{corr_file}\n"
                "请先运行graph_builder.py生成该文件（需完成图模型构建步骤）"
            )

        # 读取数据并验证字段
        corr_df = pd.read_csv(corr_file)
        required_cols = ["industry1", "industry2", "correlation"]
        missing_cols = [col for col in required_cols if col not in corr_df.columns]
        if missing_cols:
            raise ValueError(f"❌ 行业相关度数据缺少必要字段：{missing_cols}")

        # 构建双向相关度映射（行业1-行业2与行业2-行业1相关度相同）
        for _, row in corr_df.iterrows():
            ind1 = row["industry1"]
            ind2 = row["industry2"]
            corr_score = row["correlation"]
            self.industry_corr[(ind1, ind2)] = corr_score
            self.industry_corr[(ind2, ind1)] = corr_score

    def _get_all_industries(self):
        """从相关度数据中提取所有行业（去重）"""
        industries = set()
        for ind1, ind2 in self.industry_corr.keys():
            industries.add(ind1)
            industries.add(ind2)
        self.all_industries = sorted(list(industries))

    # ---------------- 辅助方法 ----------------
    def _find_similar_industries(self, input_ind: str) -> List[str]:
        """输入行业不存在时，推荐相似名称的行业"""
        return [ind for ind in self.all_industries
                if input_ind in ind or ind in input_ind]

    def _calculate_correlations(self, target_ind: str) -> List[Tuple[str, float]]:
        """计算目标行业与其他所有行业的相关度"""
        correlations = []
        for ind in self.all_industries:
            if ind == target_ind:
                continue  # 跳过自身
            # 获取相关度数据
            corr = self.industry_corr.get((target_ind, ind), 0.0)
            
            # 如果相关度为0，尝试从反向查找
            if corr == 0.0:
                corr = self.industry_corr.get((ind, target_ind), 0.0)
            
            correlations.append((ind, corr))
        return correlations

    # ---------------- 推荐核心逻辑 ----------------
    def recommend_expansion(self, target_ind: str, top_k: int = 5) -> Dict:
        """
        推荐相关行业和对冲行业（改进版）
        :param target_ind: 输入行业名称
        :param top_k: 每个类别推荐数量（默认5个）
        :return: 包含相关行业、对冲行业的结果字典
        """
        # 1. 验证输入行业是否存在
        if target_ind not in self.all_industries:
            similar_inds = self._find_similar_industries(target_ind)
            error_msg = f"❌ 未找到行业「{target_ind}」"
            if similar_inds:
                error_msg += f"\n可能的行业名称：{similar_inds}"
            return {"error": error_msg}

        # 2. 计算目标行业与其他行业的相关度
        all_correlations = self._calculate_correlations(target_ind)
        if not all_correlations:
            return {"error": f"❌ 无「{target_ind}」的相关行业数据"}

        # 3. 改进的排序逻辑
        # 相关行业：选择正相关度最高的行业（相关度 > 0.3）
        positive_correlations = [(ind, corr) for ind, corr in all_correlations if corr > 0.3]
        positive_correlations.sort(key=lambda x: x[1], reverse=True)
        related_inds = positive_correlations[:top_k]
        
        # 对冲行业：选择负相关度最低的行业（相关度 < -0.3）或相关度接近0的行业
        negative_correlations = [(ind, corr) for ind, corr in all_correlations if corr < -0.3]
        low_correlations = [(ind, corr) for ind, corr in all_correlations if abs(corr) < 0.1]
        
        # 优先选择负相关，其次选择低相关
        hedge_candidates = negative_correlations + low_correlations
        
        # 如果对冲候选不足，从所有行业中选择相关度最低的
        if len(hedge_candidates) < top_k:
            # 排除目标行业本身和已经选择的相关行业
            related_industry_names = {ind for ind, _ in related_inds}
            remaining_candidates = [(ind, corr) for ind, corr in all_correlations 
                                  if ind != target_ind and ind not in related_industry_names]
            hedge_candidates.extend(remaining_candidates)
        
        # 确保对冲行业有足够的相关度数据
        if not hedge_candidates:
            # 如果仍然没有候选，使用所有行业（除了目标行业和相关行业）
            related_industry_names = {ind for ind, _ in related_inds}
            hedge_candidates = [(ind, corr) for ind, corr in all_correlations 
                              if ind != target_ind and ind not in related_industry_names]
        
        # 按相关性绝对值排序，优先选择相关性最低的作为对冲
        hedge_candidates.sort(key=lambda x: abs(x[1]))
        
        # 确保返回足够数量的对冲行业
        hedge_inds = []
        seen_hedge = set()
        for ind, corr in hedge_candidates:
            if len(hedge_inds) >= top_k:
                break
            if ind not in seen_hedge:
                hedge_inds.append((ind, corr))
                seen_hedge.add(ind)
        
        # 如果仍然不足，从所有剩余行业中选择
        if len(hedge_inds) < top_k:
            remaining_all = [(ind, corr) for ind, corr in all_correlations 
                           if ind != target_ind and ind not in seen_hedge]
            for ind, corr in remaining_all:
                if len(hedge_inds) >= top_k:
                    break
                hedge_inds.append((ind, corr))
        
        # 确保返回指定数量的对冲行业，优先选择有相关度的行业
        final_hedge_inds = []
        seen_hedge_industries = set()  # 用于去重
        
        # 首先添加有相关度的行业（非0.0）
        for ind, corr in hedge_inds:
            if len(final_hedge_inds) >= top_k:
                break
            if ind in seen_hedge_industries:
                continue
            if corr != 0.0:  # 优先选择非0.0相关度的行业
                final_hedge_inds.append((ind, corr))
                seen_hedge_industries.add(ind)
        
        # 如果数量不足，添加0.0相关度的行业
        if len(final_hedge_inds) < top_k:
            for ind, corr in hedge_inds:
                if len(final_hedge_inds) >= top_k:
                    break
                if ind in seen_hedge_industries:
                    continue
                if corr == 0.0:  # 添加0.0相关度的行业
                    final_hedge_inds.append((ind, corr))
                    seen_hedge_industries.add(ind)
        
        # 如果仍然不足，从剩余行业中选择
        if len(final_hedge_inds) < top_k:
            remaining_candidates = [(i, c) for i, c in all_correlations 
                                  if i != target_ind and i not in seen_hedge_industries]
            for ind, corr in remaining_candidates:
                if len(final_hedge_inds) >= top_k:
                    break
                final_hedge_inds.append((ind, corr))
                seen_hedge_industries.add(ind)
        
        hedge_inds = final_hedge_inds[:top_k]

        # 4. 整理推荐结果
        return {
            "related": [
                {
                    "industry": ind,
                    "correlation": round(corr, 4),
                    "reason": self._get_related_reason(target_ind, ind, corr)
                } for ind, corr in related_inds
            ],
            "hedge": [
                {
                    "industry": ind,
                    "correlation": round(corr, 4),
                    "reason": self._get_hedge_reason(target_ind, ind, corr)
                } for ind, corr in hedge_inds
            ]
        }

    def _get_related_reason(self, target_ind: str, related_ind: str, corr: float) -> str:
        """生成相关行业的推荐理由"""
        if corr > 0.7:
            return f"与「{target_ind}」高度正相关（{round(corr, 4)}），行业走势高度同步，适合联动投资"
        elif corr > 0.5:
            return f"与「{target_ind}」中度正相关（{round(corr, 4)}），行业趋势相似，适合协同布局"
        else:
            return f"与「{target_ind}」轻度正相关（{round(corr, 4)}），有一定联动性，可考虑组合投资"

    def _get_hedge_reason(self, target_ind: str, hedge_ind: str, corr: float) -> str:
        """生成对冲行业的推荐理由"""
        if corr < -0.5:
            return f"与「{target_ind}」负相关（{round(corr, 4)}），行业走势相反，适合对冲风险"
        elif corr < -0.3:
            return f"与「{target_ind}」轻度负相关（{round(corr, 4)}），有一定对冲效果"
        else:
            return f"与「{target_ind}」相关性极低（{round(corr, 4)}），行业独立性強，适合分散投资"

    # ---------------- 可视化功能 ----------------
    def visualize_industry_relations(self, target_ind: str, result: Dict):
        """可视化目标行业与相关/对冲行业的关系"""
        # 1. 准备要展示的行业节点
        related_industries = [item["industry"] for item in result["related"]]
        hedge_industries = [item["industry"] for item in result["hedge"]]
        all_nodes = [target_ind] + related_industries + hedge_industries
        
        # 2. 创建图并添加节点
        G = nx.Graph()
        G.add_nodes_from(all_nodes)
        
        # 3. 添加边（目标行业与其他行业的关联）
        edges = []
        for ind in related_industries:
            corr = self.industry_corr.get((target_ind, ind), 0.0)
            edges.append((target_ind, ind, corr))
        
        for ind in hedge_industries:
            corr = self.industry_corr.get((target_ind, ind), 0.0)
            edges.append((target_ind, ind, corr))
        
        # 添加边到图中
        for u, v, corr in edges:
            G.add_edge(u, v, weight=corr)
        
        # 4. 配置中文字体
        plt.rcParams.update({
            "font.family": self.chinese_font,
            "axes.unicode_minus": False,
            "figure.dpi": 150
        })
        
        # 5. 设置节点样式
        fig, ax = plt.subplots(figsize=(12, 8), facecolor="white")
        
        # 节点布局
        pos = nx.spring_layout(G, seed=42, k=0.8)  # 固定seed确保布局一致
        
        # 节点颜色和大小
        node_colors = []
        node_sizes = []
        for node in G.nodes:
            if node == target_ind:
                node_colors.append("#FF6B6B")  # 目标行业：红色
                node_sizes.append(3000)
            elif node in related_industries:
                node_colors.append("#4ECDC4")  # 相关行业：青色
                node_sizes.append(2000)
            else:  # 对冲行业
                node_colors.append("#FFD166")  # 对冲行业：黄色
                node_sizes.append(2000)
        
        # 6. 绘制节点和边
        nx.draw_networkx_nodes(
            G, pos, ax=ax,
            node_color=node_colors,
            node_size=node_sizes,
            edgecolors="#333333",
            linewidths=1.0
        )
        
        # 绘制边，根据相关度调整线宽
        for u, v, data in G.edges(data=True):
            width = 2.0 + abs(data["weight"]) * 3  # 相关度越高线越粗
            nx.draw_networkx_edges(
                G, pos, ax=ax,
                edgelist=[(u, v)],
                width=width,
                edge_color="#7F7F7F"
            )
        
        # 7. 添加节点标签
        nx.draw_networkx_labels(
            G, pos, ax=ax,
            font_family=self.chinese_font,
            font_size=10,
            font_weight="bold"
        )
        
        # 8. 添加边标签（相关度）
        edge_labels = {(u, v): f"{d['weight']:.4f}" for u, v, d in G.edges(data=True)}
        nx.draw_networkx_edge_labels(
            G, pos, ax=ax,
            edge_labels=edge_labels,
            font_family=self.chinese_font,
            font_size=8,
            font_color="#DC143C",
            bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8)
        )
        
        # 9. 设置标题和保存
        plt.title(
            f"{target_ind}行业的相关与对冲行业关系图",
            fontsize=14,
            pad=20
        )
        
        # 添加图例
        import matplotlib.patches as mpatches
        red_patch = mpatches.Patch(color="#FF6B6B", label="目标行业")
        blue_patch = mpatches.Patch(color="#4ECDC4", label="相关行业")
        yellow_patch = mpatches.Patch(color="#FFD166", label="对冲行业")
        plt.legend(handles=[red_patch, blue_patch, yellow_patch], loc="best")
        
        # 保存图像
        save_name = f"industry_{target_ind.replace('/', '_')}_relations.png"
        save_path = os.path.join(self.data_dir, save_name)
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches="tight")
        plt.close()
        print(f"✅ 关系图已保存至：{os.path.abspath(save_path)}")

    # ---------------- 交互式模式 ----------------
    def interactive_mode(self):
        """交互式推荐模式：支持用户输入并实时显示结果，提供可视化选项"""
        while True:
            user_input = input("请输入要拓展的行业名称：").strip()
            if not user_input:
                print("❌ 输入不能为空，请重新输入\n")
                continue

            # 获取推荐结果
            result = self.recommend_expansion(user_input)
            if "error" in result:
                print(result["error"] + "\n")
                continue

            # 显示结果
            print(f"\n📊 「{user_input}」行业拓展推荐结果：")

            # 相关行业
            print("\n【相关行业（适合深化布局）】：")
            for i, item in enumerate(result["related"], 1):
                print(f"{i}. 行业名称：{item['industry']}")
                print(f"   相关度：{item['correlation']}")
                print(f"   推荐理由：{item['reason']}\n")

            # 对冲行业
            print("【对冲行业（适合风险分散）】：")
            for i, item in enumerate(result["hedge"], 1):
                print(f"{i}. 行业名称：{item['industry']}")
                print(f"   相关度：{item['correlation']}")
                print(f"   推荐理由：{item['reason']}\n")

            # 询问是否可视化
            viz_choice = input("是否生成行业关系图？（输入y生成，其他键跳过）：").strip().lower()
            if viz_choice == "y":
                self.visualize_industry_relations(user_input, result)

            # 询问是否继续
            continue_choice = input("是否继续查询其他行业？（输入y继续，其他键退出）：").strip().lower()
            if continue_choice != "y":
                print("👋 感谢使用，程序退出！")
                break
            print("\n" + "-" * 60 + "\n")


if __name__ == "__main__":
    try:
        # 初始化推荐器（数据目录默认"hs300_data"，若路径不同可修改）
        recommender = IndustryExpansionRecommender(data_dir="hs300_data")
        # 启动交互式模式
        recommender.interactive_mode()
    except Exception as e:
        print(f"程序运行出错：{e}")
